package com.gin.utils

import com.vesoft.nebula.client.graph.NebulaPoolConfig
import com.vesoft.nebula.client.graph.data.{HostAddress, ResultSet}
import com.vesoft.nebula.client.graph.net.{NebulaPool, Session}
import org.apache.log4j.Logger

import scala.collection.JavaConverters._
import scala.collection.mutable.ListBuffer

/**
 * GraphProvider for Nebula Graph Service
 */
class NebulaClientUdf(nebulaIps: String, timeout: Int)
  extends AutoCloseable
    with Serializable {
  private[this] lazy val LOG = Logger.getLogger(this.getClass)

  @transient val nebulaPoolConfig = new NebulaPoolConfig
  @transient val pool: NebulaPool = new NebulaPool

  var address = new ListBuffer[HostAddress]()
  if (nebulaIps.contains(",")) {
    for (ip <- nebulaIps.split(",")) {
      address.+=(new HostAddress(ip, 9669))
    }
  } else {
    address = address.+=(new HostAddress(nebulaIps, 9669))
  }

  //随机一个连接
  val randAddr = scala.util.Random.shuffle(address)
  nebulaPoolConfig.setTimeout(timeout)
  pool.init(randAddr.asJava, nebulaPoolConfig)


  def getGraphClient(userName: String, pwd: String): Session = {
    pool.getSession(userName, pwd, true);
  }

  def releaseGraphClient(session: Session): Unit = {
    session.release()
  }

  override def close(): Unit = {
    pool.close()
  }

  def switchSpace(session: Session, space: String): ResultSet = {
    val switchStatment = s"use $space"
    LOG.info(s"switch space $space")
    val result = submit(session, switchStatment)
    result
  }

  def submit(session: Session, statement: String): ResultSet = {
    session.execute(statement)
  }
}
